import os

from langchain_community.embeddings import OllamaEmbeddings

os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'  # 国内可用
os.environ['HF_HUB_DISABLE_SYMLINKS'] = '1'  # 避免 Windows 符号链接报错

from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer

SINGLE_SENT = "今天天气真好"
MULTI_SENTS = ["今天天气真好", "今天阳光不错", "股票大涨"]


def vectorize():
    model = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")  # 首次自动下载

    emb = model.encode(SINGLE_SENT)
    print(emb.shape)  # (384,)

    embs = model.encode(MULTI_SENTS)

    # 余弦相似度矩阵
    print(cosine_similarity(embs))


def vectorize_by_ollama():
    model = OllamaEmbeddings(
        model="bge-m3",
        base_url="http://localhost:11434"
    )

    vec = model.embed_query(SINGLE_SENT)  # 单条
    print(len(vec))

    embs = model.embed_documents(MULTI_SENTS)

    # 余弦相似度矩阵
    print(cosine_similarity(embs))

    embs = model.embed_documents(["""# 总体说明

javalang是一个解析java8语法的python三方库，采用的是自顶向下的解析方法，提前看1~2个字符+回溯。



tree.py   放ast的各种Node结构

parser.py    解析java8语法的核心类。

tokenizer.py    lexer



## parser详解


### lookahead函数

```
next: 从lookahead表里取出一个token消费
look：看lookahead表的下一个token，不取出
```



### accept函数

```
would_accept: 只判断是否指定token，不实际消费token
try_accept: 判断+成功则消费token+失败无动作
accept： 必须消费指定token，失败则抛JavaSyntaxError异常
```

accept既能接受指定值，也能接受指定类型，且可一次消费多个值或类型。其底层调用lookahead接口。






### 表达式解析""", "python三方包javalang的accept函数说明"])
    print(cosine_similarity(embs))


if __name__ == '__main__':
    vectorize_by_ollama()
